import base64
from io import BytesIO

import tqdm
import numpy as np
import pandas as pd
from PIL import Image
from insightface.app import FaceAnalysis



def detect_faces(app, imgs):
    n_faces = []
    for img_b64 in tqdm.tqdm(imgs.img_b64):
        img = Image.open(BytesIO(base64.b64decode(img_b64)))
        img = img.convert("RGB")
        img = np.asarray(img)

        n_faces.append(len(app.get(img)))

    return n_faces


def main(
    fp_imgs: str,
    fp_out: str,
    conf_thresh: float = 0.7,
):
    imgs = pd.read_csv(fp_imgs)
    app = FaceAnalysis(allowed_modules=["detection"])
    app.prepare(ctx_id=0, det_size=(640, 640), det_thresh=conf_thresh)

    # detect faces
    imgs["n_faces"] = detect_faces(app, imgs)

    cols = ["qry_id", "sengine", "img_rank", "n_people", "n_faces"]
    imgs[cols].to_csv(fp_out, index=False)


if __name__ == "__main__":
    main()
